package com.spring.ai.config;

import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.rag.content.retriever.ContentRetriever;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.store.embedding.pinecone.PineconeServerlessIndexConfig;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.pinecone.PineconeEmbeddingStore;

/**
 *  创建向量数据库和检索向量数据 
 */
@Configuration
public class EmbeddingStoreConfig {
    // 向量模型
    @Autowired
    private EmbeddingModel embeddingModel;
    
    // 从配置文件读取 Pinecone API key
    @Value("${pinecone.api-key}")
    private String pineconeApiKey;

    @Value("${pinecone.index-name}")
    private String pineconeIndexName;

    @Value("${pinecone.name-space}")
    private String pineconeNamespace;

    
    /**
     * 创建向量数据库
     * 使用内存存储作为示例，你可以根据需要替换为 Pinecone 或其他向量数据库
     */
    @Bean
    public EmbeddingStore<TextSegment> embeddingStore() {
        //创建向量存储
        EmbeddingStore<TextSegment> embeddingStore = PineconeEmbeddingStore.builder()
                .apiKey(pineconeApiKey)
                .index(pineconeIndexName)//如果指定的索引不存在，将创建一个新的索引
                .nameSpace(pineconeNamespace) //如果指定的名称空间不存在，将创建一个新的名称
                .createIndex(PineconeServerlessIndexConfig.builder()
                                .cloud("AWS") //指定索引部署在 AWS 云服务上。
                                .region("us-east-1") //指定索引所在的 AWS 区域为 us-east-1。
                                .dimension(embeddingModel.dimension()) //指定索引的向量维度，该维度
                        .build())
                .build();
        return embeddingStore;
    }

    /**
     * 检索向量数据
     */
    @Bean(name = "contentRetrieverEmbedding")
    ContentRetriever  contentRetrieverEmbedding(EmbeddingStore<TextSegment> embeddingStore) {
        return EmbeddingStoreContentRetriever.builder()
                .embeddingModel(embeddingModel)
                .embeddingStore(embeddingStore)
                .maxResults(2)  // 2个结果
                .minScore(0.7)  // 0.7分
                .build();
    }

}
